import torch
import torch.nn as nn
import torch.nn.functional as F

class Generator(nn.Module):
    def __init__(self, latent_dim, signal_length, n_classes=10):
        super(Generator, self).__init__()
        self.n_classes = n_classes
        
        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(latent_dim + n_classes, 512, normalize=False),
            *block(512, 1024),
            *block(1024, 2048),
            *block(2048, 4096),
            *block(4096, 2048),
            *block(2048, 1024),
            nn.Linear(1024, signal_length),
            nn.Tanh()
        )

    def forward(self, z, labels):
        # 将标签转换为one-hot编码
        labels = labels.long()
        one_hot = torch.zeros(labels.size(0), self.n_classes, device=labels.device)
        one_hot.scatter_(1, labels.unsqueeze(1), 1)
        
        # 将噪声和one-hot标签连接
        gen_input = torch.cat((one_hot, z), -1)
        signal = self.model(gen_input)
        return signal

class Discriminator(nn.Module):
    def __init__(self, signal_length, n_classes=10):
        super(Discriminator, self).__init__()
        self.n_classes = n_classes
        
        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(signal_length + n_classes, 1024, normalize=False),
            *block(1024, 512),
            *block(512, 256),
            *block(256, 128),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, signal, labels):
        # 将标签转换为one-hot编码
        labels = labels.long()
        one_hot = torch.zeros(labels.size(0), self.n_classes, device=labels.device)
        one_hot.scatter_(1, labels.unsqueeze(1), 1)
        
        # 将信号和one-hot标签连接
        d_in = torch.cat((signal, one_hot), -1)
        validity = self.model(d_in)
        return validity

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        if m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
    elif classname.find('Embedding') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)